import torch.nn as nn

from .attention import MultiHeadedAttention
from .utils import SublayerConnection,PositionwiseFeedForward

class TransformerBlock(nn.Module):
    def __init__(self,hidden,attn_heads,feed_forward_hidden,dropout):
        super().__init__()
        self.attention=MultiHeadedAttention(h=attn_heads,d_model=hidden,dropout=dropout)
        self.feed_forward=PositionwiseFeedForward(d_model=hidden,d__ff=feed_forward_hidden,dropout=dropout)
        self.input_sublayer=SublayerConnection(size=hidden,dropout=dropout)
        self.output_sublayer=SublayerConnection(size=hidden,dropout=dropout)
        self.dropout=nn.Dropout(p=dropout)
    
    def forward(self,x,mask):
        x=self.input_sublayer(x,lambda _x:self.attention.forward(_x,_x,_x,mask=mask))
        x=self.output_sublayer(x,self.feed_forward)

        return self.dropout(x)